#  The file created by Simon J. Bale
#  Modified by Yuan W.
#  Awesomized by David L.
#  DANGER: Attempts to read ALL .dat files in working directory
#!/usr/bin/python

import pylab
import sys
import re
#import cStringIO
import time
import numpy as np
from matplotlib import pyplot

import cPickle as pickle

import os, errno

def main(*args):
    try:
    	# Perhaps we could have a switch to say that we want to plot one datafile
    	# inparticular rather than all of them. This would mean we could keep
    	# aiming for a universal script
    	
    	
        startTime = time.clock()
        plotDataList = []
        
        # ****************** #
        # PROCESS ARGUEMENTS #
        # ****************** #
        drawTrans = True
	drawFreq = True
        drawMeas = True
        saveTransData = True
	saveFreqData = True
	saveMeasData = True
        fromFile = False
        requestedNumberOfFiles = 0
        noDisplay = False
        
        if "no_trans" in args:
        	drawTrans = False
        	
        if "no_meas" in args:
        	drawMeas = False
	
	if "no_freq" in args:
		drawFreq = False
        	
        if "1" in args:
        	drawMeas = False
        	requestedNumberOfFiles = 1
        	
        if "no_plot" in args:
        	drawMeas = False
        	drawTrans = False
		drawFreq = False
        	
        if "no_save_data" in args:
        	saveTransData = False
		saveFreqData = False
		saveMeasData = False
        	
        if "process_only" in args:
        	saveTransData = True
        	saveFreqData = True
        	saveMeasData = True
        	drawMeas = False
        	drawTrans = False
		drawFreq = False
        	
        if "from_file" in args:
        	fromFile = True
        	saveTransData = False
		saveFreqData = False
		saveMeasData = False
        
        if "no_display" in args:
        	noDisplay = True
        
        # ******* #
        # GOTCHAS #
        # ******* #
        if ("1" in args) and ("no_trans" in args):
        	print "I refuse to plot one measurment result (is that really what you want?)."
        	exit()
        
#        dataFile = open(args[1], 'r')

	allFilesList = os.listdir("./")
	fileList = []
	numOfDatFiles = 0
	# For ordering
	for filename in allFilesList:
		fileNumber = re.findall(r"[0-9]+", filename)
		if (filename[-4:] == ".dat" and fileNumber):
			fileList.append(0)
	for filename in allFilesList:
		fileNumber = re.findall(r"[0-9]+", filename)
		if (filename[-4:] == ".dat" and fileNumber):
			fileList[(int(fileNumber[len(fileNumber)-1])-1)] = filename
	
        indexReached = False
        endOfData = False
        
        rawTransData = list()
	rawFreqData = list()
        measurementData = []
        measurementDataHeaders = []
        
        transDataHeaders = list()
        transDataSetsCount = 0
        transData = list()

	freqDataHeaders = list()
        freqDataSetsCount = 0
        freqData = list()
        
        # For safety, see loop below
        inTransDataSet = False
	inFreqDataSet = False

	if not fromFile:
		if requestedNumberOfFiles == 0:
			fileRange = range(len(fileList))
		else:
			fileRange = range(1)
		
		for i in fileRange:
			rawTransData.append(list())
			rawFreqData.append(list())
			transDataSetsCount = 0
			freqDataSetsCount = 0
			measureStatements = 0
			with open(fileList[i], 'r') as dataFile:
				for line in dataFile:
					# These I've just used for convenience as they should recognise data and headers correctly
					# in all cases, but it would be good to improve them to be more specific
					#indexTimeLine = re.findall(r"^Index\s+(time|v-sweep)\s+.+$", line, flags=re.IGNORECASE)
					dataLine = re.findall(r"^\d+\t.*\t.*\s*$", line)
			
					# Modified line to only get time based results (hopefully)
					indexTimeLine = re.findall(r"^Index\s+(time)\s+.+$", line, flags=re.IGNORECASE)
					indexFrequencyLine = re.findall(r"^Index\s+(frequency)\s+.+$", line, flags=re.IGNORECASE)
					
					blankLine = re.findall(r"\A\s*\Z", line)
					
					measureLine = re.findall(r"^.*(targ=).*(trig=).*", line)
			
					if indexTimeLine and (saveTransData or drawTrans):
						transDataSetsCount = transDataSetsCount + 1
						rawTransData[i].append(list())			
						for element in line.split():
							# This seems pretty dirty, feel free to improve it
							if not((element == "Index") or (element == "time")):
								transDataHeaders.append(element)
						inTransDataSet = True

					if indexFrequencyLine and (saveFreqData or drawFreq):
						freqDataSetsCount = freqDataSetsCount + 1
						rawFreqData[i].append(list())
						for element in line.split():
							if not((element == "Index") or (element == "frequency")):
								freqDataHeaders.append(element)
						inFreqDataSet = True
						
					elif dataLine and inTransDataSet:
						rawTransData[i][(transDataSetsCount - 1)].append(line.split())

					elif dataLine and inFreqDataSet:
						rawFreqData[i][(freqDataSetsCount - 1)].append(line.split())

					
					elif measureLine:
						if (i == 0):
							measurementData.append(list())
							measurementDataHeaders.append(line.split()[0])
											
						measurementData[measureStatements].append(line.split()[2])
						measureStatements = measureStatements + 1
					
					# To test	
					elif blankLine:
						inTransDataSet = False
						inFreqDataSet = False
				
	if not fromFile:
	# *********************************** #
	#  TRANSIENT DATA PROCESSING SECTION  #
	# *********************************** #
		if (len(rawTransData[0]) > 0):
			# Count total number of columns in all sets, minus index and time
			# (might mess up with v-sweeps)
			totalNumOfTransDataColumns = 0
			specificNumTransDataColumns = []
			# This might be useful
			for i in range(len(rawTransData[0])):
				totalNumOfTransDataColumns = totalNumOfTransDataColumns + (len(rawTransData[0][i][0]) - 2)
				specificNumTransDataColumns.append(len(rawTransData[0][i][0]) - 2)
			
			specificTransDataColumnBeginnings = [0]*len(specificNumTransDataColumns) #totalNumOfTransDataColumns
			for i in range(len(specificNumTransDataColumns)):		
				for j in range(len(specificNumTransDataColumns)):
					if ( i == 0 ):
						specificTransDataColumnBeginnings[i] = 0
						break
					else:
						if j == i:
							break
						else:
							specificTransDataColumnBeginnings[i] = specificTransDataColumnBeginnings[i] + specificNumTransDataColumns[j]				
		
		
			# Make transData the right size
			# ( "+2" for the Index and Time columns)
			for i in range(len(rawTransData)):
				transData.append(list())
				for j in range(totalNumOfTransDataColumns + 2):
					transData[i].append(list())
					for k in range(len(rawTransData[i][0])):				
						transData[i][j].append(list())
		
			# For each file
			for i in range(len(rawTransData)):
				# For each data set
				for j in range(len(rawTransData[i])):
					# For each row in the data set
					for k in range(len(rawTransData[i][j])):
						# Take Index and Time values from first dataSet
						if j == 0:
							# Append the first element of that row to the transData[0]
							transData[i][0][k] = rawTransData[i][j][k][0]
							
							# And the second element to transData[1]
							transData[i][1][k] = rawTransData[i][j][k][1]
						
						# Go along row and copy data to transData
						for l in range(len(rawTransData[i][j][k])-2):
							transData[i][specificTransDataColumnBeginnings[j]+2+l][k] = rawTransData[i][j][k][l+2]
	
		# *********************************** #
		#  FREQUENCY DATA PROCESSING SECTION  #
		# *********************************** #
		if (len(rawFreqData[0]) > 0):
			
			# Count total number of columns in all sets, minus index and time
			# (might mess up with v-sweeps)
			totalNumOfFreqDataColumns = 0
			specificNumFreqDataColumns = []
			# This might be useful
			for i in range(len(rawFreqData[0])):
				totalNumOfFreqDataColumns = totalNumOfFreqDataColumns + (len(rawFreqData[0][i][0]) - 2)
				specificNumFreqDataColumns.append(len(rawFreqData[0][i][0]) - 2)
			
			specificFreqDataColumnBeginnings = [0]*len(specificNumFreqDataColumns) #totalNumOfTransDataColumns
			for i in range(len(specificNumFreqDataColumns)):		
				for j in range(len(specificNumFreqDataColumns)):
					if ( i == 0 ):
						specificFreqDataColumnBeginnings[i] = 0
						break
					else:
						if j == i:
							break
						else:
							specificFreqDataColumnBeginnings[i] = specificFreqDataColumnBeginnings[i] + specificFreqTransDataColumns[j]				
		
		
			# Make transData the right size
			# ( "+2" for the Index and Time columns)
			for i in range(len(rawFreqData)):
				freqData.append(list())
				for j in range(totalNumOfFreqDataColumns + 2):
					freqData[i].append(list())
					for k in range(len(rawFreqData[i][0])):				
						freqData[i][j].append(list())
		
			# For each file
			for i in range(len(rawFreqData)):
				# For each data set
				for j in range(len(rawFreqData[i])):
					# For each row in the data set
					for k in range(len(rawFreqData[i][j])):
						# Take Index and Time values from first dataSet
						if j == 0:
							# Append the first element of that row to the transData[0]
							freqData[i][0][k] = rawFreqData[i][j][k][0]
							
							# And the second element to transData[1]
							freqData[i][1][k] = rawFreqData[i][j][k][1]
						
						# Go along row and copy data to transData
						for l in range(len(rawFreqData[i][j][k])-2):
							freqData[i][specificFreqDataColumnBeginnings[j]+2+l][k] = rawFreqData[i][j][k][l+2]
	
	if fromFile:
		print "Using pre-processed data files."
		transDataExists = False
		transDataHeadersExists = False
		measurementDataHeadersDataExists = False
		measurementDataDataExists = False
		if os.path.exists("./plots/processed_data/trans_data.dat"):
			transDataExists = True
		else:
			print "No Trans data found."
		if os.path.exists("./plots/processed_data/trans_data_headers.dat"):
			transDataHeadersExists = True
		else:
			print "No Trans data headers found."
		if os.path.exists("./plots/processed_data/measure_names_data.dat"):
			measurementDataHeadersDataExists = True
		else:
			print "No Measurement data headers found"
		if os.path.exists("./plots/processed_data/measure_lines_data.dat"):
			measurementDataDataExists = True
		else:
			print "No Measurement data found"
			
		if transDataExists and transDataHeadersExists:
			transDataFile = open("./plots/processed_data/trans_data.dat","r")
			transData = pickle.load(transDataFile)
			transDataFile.close()
			transDataHeadersFile = open("./plots/processed_data/trans_data_headers.dat","r")
			transDataHeaders = pickle.load(transDataHeadersFile)
			transDataFile.close()
		
		if measurementDataHeadersDataExists and measurementDataDataExists:
			measurementDataHeadersFile = open("./plots/processed_data/measure_names_data.dat","r")
			measurementDataHeaders = pickle.load(measurementDataHeadersFile)
			measurementDataHeadersFile.close()
			measurementDataFile = open("./plots/processed_data/measure_lines_data.dat","r")
			measurementData = pickle.load(measurementDataFile)			
			measurementDataFile.close()
		elif measurementDataHeadersDataExists and not measurementDataDataExists:
			print "Measurement names exists but not data" 
		elif measurementDataDataExists and not measurementDataHeadersDataExists:
			print "Measurement data exists but not names"
	

	# ************** #
	# GRAPH SECTION  #
	# ************** #
	if (transData or measurementDataHeaders or freqData) and ((drawTrans or drawMeas or drawFreq) or saveTransData):
		if not os.path.exists('./plots/'):
			os.makedirs('./plots/')
	else:
		exit()
		
	# ***************** #
	# SAVE DATA SECTION #
	# ***************** #
	if saveTransData:
		if not os.path.exists('./plots/processed_data/'):
			os.makedirs('./plots/processed_data/')
			
		if transData:
			transDataFile = open("./plots/processed_data/trans_data.dat","w")
			pickle.dump(transData, transDataFile)
			transDataFile.close()
			transDataHeadersFile = open("./plots/processed_data/trans_data_headers.dat","w")
			pickle.dump(transDataHeaders, transDataHeadersFile)
			transDataFile.close()
		if measurementDataHeaders:
			measurementDataHeadersFile = open("./plots/processed_data/measure_names_data.dat","w")
			pickle.dump(measurementDataHeaders, measurementDataHeadersFile)
			measurementDataHeadersFile.close()
			measurementDataFile = open("./plots/processed_data/measure_lines_data.dat","w")
			pickle.dump(measurementData, measurementDataFile)
			measurementDataFile.close()
		
	# Keep track of how many figures we have
	figureCounter = 1
	
	#if (len(transData[0]) > 0) and drawTrans:
	if transData and drawTrans:
#		pyplot.figure(num=figureCounter, figsize=(16,12), dpi=120)
		pyplot.figure(num=figureCounter)

		figureCounter = figureCounter + 1
				
		if requestedNumberOfFiles == 0:
			numberOfTransLines = range(len(transData))
		else:
			numberOfTransLines = range(requestedNumberOfFiles)
				
		for i in range(len(transData[0])-2):
			pyplot.subplot(len(transData[0])-2,1,(i+1))
			
			for j in numberOfTransLines:
				if j == 0:
					pyplot.plot(transData[j][1], transData[j][i+2], label=transDataHeaders[i])
				else:
					pyplot.plot(transData[j][1], transData[j][i+2], label='_nolegend_')
				
					
			pyplot.legend()
			tempArray = np.array(transData[0][i+2]).astype(float)
			diffTemp = (np.max(tempArray) - np.min(tempArray))
			maxVal = np.max(tempArray) + (diffTemp * 0.2)
			minVal = np.min(tempArray) - (diffTemp * 0.2)
			## Hardcoding 
			#minVal = -0.2
			#maxVal = 1.2
			pyplot.ylim(ymax=maxVal, ymin=minVal)
			
			# Outputs
			pylab.savefig('./plots/Plot_Trans.pdf')
			print "Trans fig. %i complete!" % (i)
	else:
		print "Not drawing trans or no data"
	
	if freqData and drawFreq:
		pyplot.figure(num=figureCounter, figsize=(16,12), dpi=120)
		figureCounter = figureCounter + 1
				
		if requestedNumberOfFiles == 0:
			numberOfFreqLines = range(len(freqData))
		else:
			numberOfFreqLines = range(requestedNumberOfFiles)
				
		for i in range(len(freqData[0])-2):
			pyplot.subplot(len(freqData[0])-2,1,(i+1))
			
			for j in numberOfFreqLines:
				if j == 0:
					pyplot.plot(freqData[j][1], freqData[j][i+2], label=freqDataHeaders[i])
				else:
					pyplot.plot(freqData[j][1], freqData[j][i+2], label='_nolegend_')
				
					
			pyplot.legend()
			tempArray = np.array(freqData[0][i+2]).astype(float)
			diffTemp = (np.max(tempArray) - np.min(tempArray))
			maxVal = np.max(tempArray) + (diffTemp * 0.2)
			minVal = np.min(tempArray) - (diffTemp * 0.2)
			## Hardcoding 
			#minVal = -0.2
			#maxVal = 1.2
			pyplot.ylim(ymax=maxVal, ymin=minVal)
			
			pyplot.xscale('log')
			
			# Outputs
			pylab.savefig('./plots/Plot_Freq.pdf')
			print "Freq fig. %i complete!" % (i)
	else:
		print "Not drawing freq or no data"

	if measurementDataHeaders and measurementData and drawMeas:
		for i in range(len(measurementDataHeaders)):
			#pyplot.figure(num=figureCounter, figsize=(12,8), dpi=300)
			pyplot.figure(num=figureCounter, figsize=(12,8))
			figureCounter = figureCounter + 1
			
			# Boxplot
			pyplot.subplot(121)
			pyplot.title(measurementDataHeaders[i])			
			pyplot.ylabel('Time')
			x = np.array(measurementData[i])
			pyplot.boxplot(x.astype(float))
			
			# VLine plot
			pyplot.subplot(122)
			pyplot.title(measurementDataHeaders[i])
			pyplot.xlabel('Circuit Number')
			pyplot.ylabel('Time')
			x = np.array(measurementData[i]).astype(float)
			yMax = np.max(x) + ((np.max(x) - np.min(x)) * 0.05)
			yMin = np.min(x) - ((np.max(x) - np.min(x)) * 0.05)
			#pyplot.vlines(range(len(x)), yMin, x, label='Delay (s)', lw=2)
			circuitNumbers = []
			for counter in range(len(measurementData[i])):
				circuitNumbers.append(counter+1)
				
			pyplot.scatter(circuitNumbers, x, s=20, c='b', marker='o', cmap=None, norm=None, vmin=None, vmax=None, alpha=None, linewidths=None, verts=None, label='Delay (s)')
			
			pyplot.xlim(xmin=1, xmax=len(measurementData[i]))
			pyplot.ylim(ymin=yMin, ymax=yMax)
			
			# Outputs
			pyplot.savefig('./plots/measure-%s.pdf' % (measurementDataHeaders[i]))
			print "Meas fig. %i complete!" % (i)

	
    except:
        print "Unexpected error:", sys.exc_info()[0]
        raise
    else:
        print (time.clock() - startTime)
        if not noDisplay:
        	pyplot.show()
        return 0

if __name__ == '__main__':
    sys.exit(main(*sys.argv))
